from numpy import array,argmin,dot,cross,zeros
from v3dsolver.fluid_elements.fluid_elems import Saffman
from v3dsolver.vector import Vec

def centered3(x,reverse=False):
    """
    A third order centered interpolation of an element located at 
    -1/2<=x<=1/2. The output is a list of tuples each of which contain
    the position of the point of interpolation and the fraction awarded
    to it. 
    
    reverse: tells whether distribution is to be done in the negative 
             direction (default: False)
    """
    if reverse:
        x = -x
    result = [[-3./2,(3-2*x)*(4*x**2-1)/48.],
              [-1./2,(1-2*x)*(9-4*x**2)/16.],
              [ 1./2,(1+2*x)*(9-4*x**2)/16.],
              [ 3./2,(3+2*x)*(4*x**2-1)/48.]]
    if reverse:
        for i in range(len(result)):
            result[i][0] = - result[i][0]
    return result
    
def decentered3(x,reverse=False):
    """
    A third order decentered interpolation of an element located at 
    -1/2<=x<=1/2. The output is a list of tuples each of which contain
    the position of the point of interpolation and the fraction awarded
    to it. 
    
    reverse: tells whether distribution is to be done in the negative 
             direction (default: False)
    """
    if reverse:
        x = -x
    result = [[-1./2,(1-2*x)*(2*x-5)*(2*x-3)/48.],
              [ 1./2,(2*x-5)*(2*x-3)*(1+2*x)/16.],
              [ 3./2,(1-2*x)*(2*x-5)*(1+2*x)/16.],
              [ 5./2,(1-2*x)*(3-2*x)*(1+2*x)/48.]]
    if reverse:
        for i in range(len(result)):
            result[i][0] = - result[i][0]
    return result
    
def centered2(x,reverse=False):
    """
    A second order centered interpolation of an element located at 
    -1/2<=x<=1/2. The output is a list of tuples each of which contain
    the position of the point of interpolation and the fraction awarded
    to it. 
    
    reverse: tells whether distribution is to be done in the negative 
             direction (default: False)
    """
    if reverse:
        x = -x
    result = [[-1,x*(x-1)/2.],
              [ 0,(1-x**2)],
              [ 1,x*(1+x)/2.]]
    if reverse:
        for i in range(len(result)):
            result[i][0] = - result[i][0]
    return result    
    
def decentered2(x,reverse=False):
    """
    A second order decentered interpolation of an element located at 
    -1/2<=x<=1/2. The output is a list of tuples each of which contain
    the position of the point of interpolation and the fraction awarded
    to it. 
    
    reverse: tells whether distribution is to be done in the negative 
             direction (default: False)
    """
    if reverse:
        x = -x
    result = [[0,(x-2)*(x-1)/2.],
              [1,x*(2-x)],
              [2,x*(x-1)/2.]]
    if reverse:
        for i in range(len(result)):
            result[i][0] = - result[i][0]
    return result
    
def centered1(x,reverse=False):
    """
    A first order centered interpolation of an element located at 
    -1/2<=x<=1/2. The output is a list of tuples each of which contain
    the position of the point of interpolation and the fraction awarded
    to it.
    
    reverse: tells whether distribution is to be done in the negative 
             direction (default: False) 
    """
    if reverse:
        x = -x
    result = [[0,(1-x)],
              [1, x]]
    if reverse:
        for i in range(len(result)):
            result[i][0] = - result[i][0]
    return result    

def nearest_pt(x,reverse=False):
    """
    An interpolation in which particle vorticity is assigned to the nearest
    grid-point to its left. The output is a list of tuples each of which 
    contain the position of the point of interpolation and the fraction 
    awarded to it.
    
    reverse: tells whether distribution is to be done in the negative 
             direction (default: False)
    """
    if reverse:
        x = -x
    result = [[-1/2.,1]]
    if reverse:
        for i in range(len(result)):
            result[i][0] = - result[i][0]
    return result

class Grid():
    """
    The grid class containing grid size, domain and panels intersecting 
    it.
    """    
    def __init__(self,domain,spacing,geometry,rad):
        """
        Initialization of Grid class.
        
        domain: domain of concern, in the form [xmin,ymin,zmin,length]
        spacing: grid spacing desired
        geometry: Geometry object containing boundary data
        rad: radius of blobs
        """
        self.domain = domain
        self.spacing = spacing
        self.geometry = geometry
        self.size = int(round(self.domain[3]/self.spacing))
        self.rad = rad
        self.initialize_grid()
    
    def initialize_grid(self):
        """
        Creates an all-zero 3D array for vorticity and another 3D array
        containing positions of grid points
        """
        self.vors_grid_pts = zeros((self.size,self.size,self.size,3),
                                   dtype='object')
        self.pos_grid_pts = zeros((self.size,self.size,self.size,3),
                                   dtype='object')
        for i in range(self.size):
            for j in range(self.size):
                for k in range(self.size):
                    self.pos_grid_pts[i,j,k] = array([self.domain[0]+i*self.spacing,
                                                     self.domain[1]+j*self.spacing,
                                                     self.domain[2]+k*self.spacing])
    
    def get_cell(self,pos):
        """
        Returns the cell id corresponding to pos in the form:
        [xindex,yindex,zindex] 
        """
        xindex = int((pos[0]-self.domain[0])/self.spacing)
        yindex = int((pos[1]-self.domain[1])/self.spacing)
        zindex = int((pos[2]-self.domain[2])/self.spacing)
        return [xindex,yindex,zindex]
    
    def get_gridpt_index(self,pos):
        """
        Returns the index of the grid point located at position pos
        in the form: [xindex,yindex,zindex]
        """
        xindex = int(round((pos[0]-self.domain[0])/self.spacing))
        yindex = int(round((pos[1]-self.domain[1])/self.spacing))
        zindex = int(round((pos[2]-self.domain[2])/self.spacing))
        return [xindex,yindex,zindex]
    
    def update_vors(self,grid_elems):
        """
        Updates the vorticities of grid points according to the list of 
        temporary grid_elems generated for each particle.
        """
        for grid_elem in grid_elems:
            xindex,yindex,zindex = self.get_gridpt_index(grid_elem.pos)
            self.vors_grid_pts[xindex,yindex,zindex] += grid_elem.vor
    
    def get_interpolated_elems(self):
        """
        Returns a list of interpolated elems located at grid points
        """
        interpolated_elems = []
        for i in range(self.size):
            for j in range(self.size):
                for k in range(self.size):
                    newelem = Saffman(pos=self.pos_grid_pts[i,j,k],
                                      vor=self.vor_grid_pts[i,j,k],
                                      vel=Vec(),rad=self.rad)
                    interpolated_elems += [newelem]
        return interpolated_elems

class SingleElemInterpolator():
    """
    Given an element, this class is reponsible for various functions used
    in distributing it's vorticity to grid-points.
    """
    def __init__(self,elem,grid):
        """
        Initialization of SingleElemInterpolator class. 
        
        elem: elem to be interpolated
        grid: Grid object containing info about grid points positions, grid
              spacing and panels intersecting each cell
        """
        self.elem = elem
        self.grid = grid
        self.parentcell = self.grid.get_cell(elem.pos)
        
    def get_interpolated_elems(self):
        """
        Performs XYZ,XZY,YXZ,YZX,ZXY,ZYX distributions and selects the
        best distribution among them according to the penalty that each
        of them provides. 
        Returns a list of final interpolated particles.
        """
        paths = [(0,1,2),(0,2,1),(1,0,2),(1,2,0),(2,0,1),(2,1,0)]
        penalties = [0,0,0,0,0,0]
        final_tmp_elems = [[],[],[],[],[],[]]
        for i,path in enumerate(paths):
            tmp_elems = [[self.elem]]
            level = 0
            total_penalty = 0
            for j in range(len(path)):
                tmp_elems += [[]]
                for elem in tmp_elems[level]: 
                    this_tmp_elems,penalty = self.distribute_vorticity(elem, 
                                                           direction=path[j])
                    tmp_elems[level+1] += this_tmp_elems
                    total_penalty += penalty
                level += 1
            final_tmp_elems[i] = tmp_elems[-1]
            penalties[i] = total_penalty
        i_min = argmin(penalties)
        return final_tmp_elems[i_min] 
                    
    def distribute_vorticity(self,elem,direction):
        """
        Distributes the vorticity of elem in the direction specified.
        
        elem: element to be distributed
        direction: an integer specifying direction (0:x,1:y,2:z)
        """
        xp = elem.pos[direction]
        x0 = (int((xp-self.grid.domain[direction])/self.grid.spacing)*
              self.grid.spacing) + self.grid.domain[direction]
        dn,dp = self.calculate_dn_dp(elem,direction)
        xdash = xp - x0
        x = xdash/self.grid.spacing - 0.5
        tmp_elems = []
        if dn > 2*self.grid.spacing:
            tmp_elems,penalty = self.case1(x,x0,elem,direction,dp)
        elif dn > self.grid.spacing:
            tmp_elems,penalty = self.case2(x,x0,elem,direction,dp)
        elif dn > 0:
            tmp_elems,penalty = self.case3(x,x0,elem,direction,dp)
        else:
            tmp_elems,penalty = self.case4(x,x0,elem,direction,dp)
        return tmp_elems,penalty
    
    def case1(self,x,x0,elem,direction,dp):
        """
        If dn > 2h. Returns list of tmp_elems and associated penalty with
        this operation.
        """
        if dp > 2*self.grid.spacing:
            tmp_elems = self.get_tmp_particles(elem,direction,centered3,x,x0)
            penalty = 0
        elif dp > self.grid.spacing:
            tmp_elems = self.get_tmp_particles(elem,direction,decentered3,x,x0,
                                               reverse=True)
            penalty = 1
        else:
            tmp_elems = self.get_tmp_particles(elem,direction,decentered2,x,x0,
                                               reverse=True,offset=0.5)
            penalty = 1
        return tmp_elems,penalty
    
    def case2(self,x,x0,elem,direction,dp):
        """
        If dn > h. Returns list of tmp_elems and associated penalty with
        this operation.
        """
        if dp > 2*self.grid.spacing:
            tmp_elems = self.get_tmp_particles(elem,direction,centered3,x,x0)
            penalty = 0
        elif dp > self.grid.spacing:
            tmp_elems = self.get_tmp_particles(elem,direction,centered2,x,x0,
                                               offset=0.5)
            penalty = 0
        else:
            tmp_elems = self.get_tmp_particles(elem,direction,centered1,x,x0,
                                               reverse=True,offset=0.5)
            penalty = 2
        return tmp_elems,penalty
    
    def case3(self,x,x0,elem,direction,dp):
        """
        If dn > 0. Returns list of tmp_elems and associated penalty with
        this operation.
        """
        if dp > 3*self.grid.spacing:
            tmp_elems = self.get_tmp_particles(elem,direction,decentered3,x,x0)
            penalty = 1
        elif dp > 2*self.grid.spacing:
            tmp_elems = self.get_tmp_particles(elem,direction,decentered2,x,x0,
                                               offset=0.5)
            penalty = 1
        elif dp > self.grid.spacing:
            tmp_elems = self.get_tmp_particles(elem,direction,centered1,x,x0,
                                               offset=0.5)
            penalty = 2
        else:
            tmp_elems = self.get_tmp_particles(elem,direction,nearest_pt,x,x0)
            penalty = 3
        return tmp_elems,penalty
    
    def case4(self,x,x0,elem,direction,dp):
        """
        If dn < 0. Returns list of tmp_elems and associated penalty with
        this operation.
        """
        if dp > 3*self.grid.spacing:
            tmp_elems = self.get_tmp_particles(elem,direction,decentered2,x,x0,
                                               offset=-0.5)
            penalty = 1
        elif dp > 2*self.grid.spacing:
            tmp_elems = self.get_tmp_particles(elem,direction,centered1,x,x0,
                                               offset=-0.5)
            penalty = 2
        elif dp > self.grid.spacing:
            tmp_elems = self.get_tmp_particles(elem,direction,nearest_pt,x,x0,
                                               reverse=True)
            penalty = 3
        else:
            tmp_elems = []
            penalty = 10
        return tmp_elems,penalty
    
    def get_tmp_particles(self,elem,direction,interpolation_kernel,x,x0,
                          reverse=False,offset=0):
        """
        Returns a list of temporary particles by distributing vorticity
        of element in the particular direction.
        
        elem: element to be interpolated
        direction: direction in which interpolation is to be done
        interpolation_kernel: interpolation kernel to be used:
                              centered3, decentered3, centered2 etc.
        x: position of element mapped between [-1/2,1/2]
        x0: position of nearest grid-point to the left of particle
        reverse: specifies whether interpolation is to be done in negative
                 direction
        offset: offset in x which will give the new x to be passed to 
                interpolation kernel 
        """
        pos_fracs = interpolation_kernel(x+offset,reverse)
        tmp_elems = []
        for pos_frac in pos_fracs:
            x_tmp,frac = pos_frac
            x_tmp = ((x_tmp-offset)+0.5)*self.grid.spacing + x0
            tmp_elem_pos = 1.0*elem.pos
            tmp_elem_pos[direction] = x_tmp
            tmp_elem_vor = frac*elem.vor
            tmp_elems += [Saffman(pos=tmp_elem_pos,vor=tmp_elem_vor,
                                  vel=Vec(),rad=elem.rad)]
        return tmp_elems
    
    def _sameside(self,p1,p2,a,b):
        """
        Returns True if points p1 and p2 lie on the same side of the line
        joining points a and b.
        """
        vec1 = (b-a).cross(p1-a)
        vec2 = (b-a).cross(p2-a)
        if vec1.dot(vec2)>=0:
            return True
        else:
            return False

    def _pt_in_triangle(self,p,a,b,c):
        """
        Returns True if point p lies inside the triangle generated by points
        a, b and c.
        """
        if (self._sameside(p,a,b,c) and self._sameside(p,b,a,c) and 
            self._sameside(p,c,a,b)):
            return True
        else:
            return False

    def calculate_dn_dp(self,elem,direction):
        """
        Calculates the nearest points where line of distribution intersects 
        any of the panels and returns dn and dp accordingly, where dn and dp
        are defined as:
        
        dn: distance of nearest pt from x0 towards left (or below) of particle 
        dp: distance of nearest pt from x0 towards right (or above) of particle
        
        elem: particle which is to be distributed
        direction: direction of distribution (0,1,2)
        """
        xp = elem.pos[direction]
        x0 = (int((xp-self.grid.domain[direction])/self.grid.spacing)*
                 self.grid.spacing)
        u = Vec()
        u[direction] = 1.
        possible_dn,possible_dp = [],[]
        for panel in self.grid.geometry.panels:
            n = panel.unitnormal
            if (not abs(n.dot(u))<1e-5):
                s1 = n.dot(panel.vertices[0]-elem.pos)/n.dot(u)
                Ps1 = elem.pos + s1*u
                if self._pt_in_triangle(Ps1,panel.vertices[0],panel.vertices[1],
                                        panel.vertices[2]):
                    if Ps1[direction]<xp: 
                        possible_dn += [x0 - Ps1[direction]]
                    if Ps1[direction]>=xp: 
                        possible_dp += [Ps1[direction] - x0]
        # if no intersection with panel is found, assign large values to dn
        # and dp
        if possible_dn==[]: 
            dn = 4*self.grid.spacing
        else:
            dn = min(possible_dn)
        if possible_dp==[]:
            dp = 4*self.grid.spacing
        else:
            dp = min(possible_dp)
        return dn,dp

class VorticityInterpolator():
    """
    This class is responsible for encompassing several functions which
    interpolate particle's vorticities to regular grid-points required
    for post-processing later.
    """
    def __init__(self,elems,domain,grid_density,geometry):
        """
        Initialization of VorticityInterpolator object.
        
        elems: list of elems local to the current processor
        grid_density: number of grid cells desired along each dimension
        geometry: a Geometry object defined in boundary_manager
        """
        self.elems = elems
        rad = self.elems[0].rad
        grid_spacing = domain[-1]/float(grid_density)
        self.grid = Grid(domain,grid_spacing,geometry,rad) 
        
    def interpolate_vorticities(self):
        """
        Distributes vorticities of elems to grid-points and adds them to
        their existing vorticities.
        
        Note: Methodology followed is exactly the same as that given in
        PR's Ph.D. thesis
        """
        for elem in self.elems:
            elem_interpolator = SingleElemInterpolator(elem,self.grid)
            interpolated_elems = elem_interpolator.get_interpolated_elems()
            self.grid.update_vors(interpolated_elems)
        return self.grid.get_interpolated_elems()
